#!/usr/bin/python -O

# homruns - a simple GUI for visualizing runs of homozygosity in SNP data
# Copyright (C) 2008 Robert Karlsson
#
# This program is distributed under the GNU General Public
# License, version 2. See the file COPYING for details.
###########################################################
#
# This program provides a graphical user interface for testing
# various parameter adjustments for the "Runs of Homozygosity"
# calculation in the plink software package for analysis of whole
# genome association data (http://pngu.mgh.harvard.edu/~purcell/plink/).
#
# The program requires a working copy of plink in the path,
# the gtk+ toolkit and its python bindings pygtk,
# and matplotlib for producing graphs.
#

import commands, os, tempfile, re, tarfile, threading, time

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as pyplot

import pygtk
pygtk.require("2.0")
import gtk
import gtk.glade

gtk.gdk.threads_init()

# plotting colors
grey60 = "#999999"
grey80 = "#CCCCCC"
grey40 = "#666666"
black = "#000000"
nicered = "#E41A1C"
nicegreen = "#4DAF4A"
niceblue = "#377EB8"
mendelorange = "#FF7F00"

# chromosome endpoints fetched from the UCSC genome browser
chrlimits = ((0,247249719), # 1
        (0,242951149),
        (0,199501827),
        (0,191273063),
        (0,180857866),
        (0,170899992),
        (0,158821424),
        (0,146274826),
        (0,140273252),
        (0,135374737),
        (0,134452384),
        (0,132349534),
        (0,114142980),
        (0,106368585),
        (0,100338915),
        (0, 88827254),
        (0, 78774742),
        (0, 76117153),
        (0, 63811651),
        (0, 62435964),
        (0, 46944323),
        (0, 49691432), # 22
        (0,154913754), # 23 == X
        (0, 57772954), # 24 == Y
        (0,154913754), # 25 == PAR, mapped on X
        (0,   16600)) # 26 == MT

class MainWindow:
    """Main User Interface class
    
    Sets up the main user interface and
    dispatches jobs to the data processing parts
    of the program as needed"""
    def __init__(self):
        """Set up"""
        # Where is this script?
        self.scriptdir = os.path.dirname(os.path.realpath(__file__))

        # Create and connect an instance of the main calculations object
        # to ourselves
        self.plink = PlinkLink()
        
        # Read and create the GUI from an XML file created in Glade.
        # Any widget we need to interact with can be fetched by name
        # using the get_widget(string) method of the wTree object.
        guifile = self.scriptdir + "/" + "homruns.glade"
        self.wTree = gtk.glade.XML(guifile)
        
        # Automatically connect glade event handlers
        # to python functions (methods belonging to this object)
        # with the same name
        self.wTree.signal_autoconnect(self)
        
        # Get references to all the widgets we somehow need
        # to interact with
        self.get_widgets()
        
        # Create a legend for the graphs
        self.setup_legend()
        
        # Create a filechooser widget
        self.setup_filechooser()

        # Load the window icon
        gtk.window_set_default_icon_from_file(self.scriptdir + "/" + "icon16x16.png")
        self.mainwindow.set_icon_from_file(self.scriptdir + "/" + "icon16x16.png")

    def get_widgets(self):
        """Get references to GUI widgets we need to interact with"""
        # the main window
        self.mainwindow = self.wTree.get_widget("mainwindow")
        # checkboxes
        self.check_density = self.wTree.get_widget("check_density")
        self.check_homozyg_kb = self.wTree.get_widget("check_homozyg_kb")
        self.check_homozyg_snp = self.wTree.get_widget("check_homozyg_snp")
        self.check_window_kb = self.wTree.get_widget("check_window_kb")
        self.check_window_snp = self.wTree.get_widget("check_window_snp")
        # window parameter spinners
        self.spin_window_kb = self.wTree.get_widget("spin_window_kb")
        self.spin_window_snp = self.wTree.get_widget("spin_window_snp")
        self.spin_window_het = self.wTree.get_widget("spin_window_het")
        self.spin_window_missing = self.wTree.get_widget("spin_window_missing")
        # segment parameter spinners
        self.spin_homozyg_kb = self.wTree.get_widget("spin_homozyg_kb")
        self.spin_homozyg_snp = self.wTree.get_widget("spin_homozyg_snp")
        self.spin_density = self.wTree.get_widget("spin_density")
        self.spin_gap = self.wTree.get_widget("spin_gap")
        # sliders
        self.hscale_threshold = self.wTree.get_widget("hscale_threshold")
        self.hscale_chrom = self.wTree.get_widget("hscale_chrom")
        # images
        self.image_chrom = self.wTree.get_widget("image_chrom")
        self.scrolledwindow_image = self.wTree.get_widget("scrolledwindow_image")
        # save button
        self.button_save = self.wTree.get_widget("button_save")
        # "please wait" frame
        self.frame_running = self.wTree.get_widget("frame_running")
        self.progressbar_plink = self.wTree.get_widget("progressbar_plink")
    
    def setup_legend(self):
        """Set up widgets for the legend panel"""
        # get references to the legend-related widgets
        self.frame_legend = self.wTree.get_widget("frame_legend")
        self.image_leg_case = self.wTree.get_widget("image_leg_case")
        self.image_leg_control = self.wTree.get_widget("image_leg_control")
        self.image_leg_unknown = self.wTree.get_widget("image_leg_unknown")
        self.image_leg_mendel = self.wTree.get_widget("image_leg_mendel")
        # initialize colors and place them in the correct legend panel spots
        self.image_leg_mendel.set_from_pixbuf(new_coloured_pixbuf(0xff7f00ff, 50, 24))
        self.image_leg_control.set_from_pixbuf(new_coloured_pixbuf(0x4daf4aff, 50, 24))
        self.image_leg_unknown.set_from_pixbuf(new_coloured_pixbuf(0x377eb8ff, 50, 24))
        self.image_leg_case.set_from_pixbuf(new_coloured_pixbuf(0xe41a1cff, 50, 24))
    
    def setup_filechooser(self):
        """Set up the file chooser"""
        self.filechooser_input = self.wTree.get_widget("filechooser_input")
        filter = gtk.FileFilter()
        filter.set_name("Pedigree files")
        filter.add_pattern("*.ped")
        filter.add_pattern("*.bed")
        self.filechooser_input.add_filter(filter)
    
    def paramcheck(self):
        """Validate filenames"""
        # checks for necessary parameters:
        # do we have a pedigree file?
        ped_filename = self.filechooser_input.get_filename()
        if ped_filename is None: raise PlinkParamError(
            "Please choose a pedigree file before running.")
        
        # is it text or binary?
        m = re.match("^(.*)\\.(p|b)ed$", ped_filename)
        self.binary_pedfile = (m.group(2) == "b")
        
        # do we have all required files?
        self.ped_basename = m.group(1)
        if self.binary_pedfile:
            look_for_file(self.ped_basename + ".bim", PlinkParamError)
            look_for_file(self.ped_basename + ".fam", PlinkParamError)
        else:
            look_for_file(self.ped_basename + ".map", PlinkParamError)
            
    
    def show_error_dlg(self, err):
        """Show an error message provided by an exception"""
        error_dlg = gtk.MessageDialog(flags = gtk.DIALOG_MODAL,
                                      type = gtk.MESSAGE_ERROR,
                                      buttons = gtk.BUTTONS_OK)
        error_dlg.set_title("Error")
        error_dlg.set_markup(str(err))
        error_dlg.run()
        error_dlg.destroy()

    def show_msg_dlg(self, msg):
        """Show an informative message"""
        msg_dlg = gtk.MessageDialog(flags = gtk.DIALOG_MODAL, 
                buttons = gtk.BUTTONS_OK, message_format = msg)
        msg_dlg.run()
        msg_dlg.destroy()

    # SIGNAL HANDLERS:
    def gtk_main_quit(self, widget):
        """Leave GTK main loop"""
        gtk.main_quit()
    
    def on_check_density_toggled(self, widget):
        """Make the density spinner follow the state of the density checkbox"""
        self.spin_density.set_sensitive(self.check_density.get_active())
    
    def on_check_homozyg_kb_toggled(self, widget):
        """Make the homozyg-kb spinner follow the state of the homozyg-kb checkbox"""
        self.spin_homozyg_kb.set_sensitive(self.check_homozyg_kb.get_active())
    
    def on_check_homozyg_snp_toggled(self, widget):
        """Make the homozyg-snp spinner follow the state of the homozyg-snp checkbox"""
        self.spin_homozyg_snp.set_sensitive(self.check_homozyg_snp.get_active())
    
    def on_button_execute_clicked(self, widget):
        """Run analysis job.
        
        Do a data sanity check.
        If that turns out OK, run main analysis + graph generation.
        If that turns out OK, show resulting graphs and legend."""
        try:
            # Check that all needed parameters are there,
            # raise error if they are not.
            self.paramcheck()
            # Fetch active parameters from GUI,
            # store them in a dictionary which we use
            # when calling our main calculations object.
            plinkdic = {}
            if self.binary_pedfile:
                plinkdic['bfile'] = self.ped_basename
            else:
                plinkdic['file'] = self.ped_basename
            if self.check_window_kb.get_active():
                plinkdic['window_kb'] = self.spin_window_kb.get_value()
            if self.check_window_snp.get_active():
                plinkdic['window_snp'] = self.spin_window_snp.get_value()
            plinkdic['window_het'] = self.spin_window_het.get_value()
            plinkdic['window_missing'] = self.spin_window_missing.get_value()
            plinkdic['window_threshold'] = self.hscale_threshold.get_value()
            if self.check_homozyg_kb.get_active():
                plinkdic['kb'] = self.spin_homozyg_kb.get_value()
            if self.check_homozyg_snp.get_active():
                plinkdic['snp'] = self.spin_homozyg_snp.get_value()
            if self.check_density.get_active():
                plinkdic['density'] = self.spin_density.get_value()
            plinkdic['gap'] = self.spin_gap.get_value()
            
            # Run main calculations
            # Let the user know something is happening
            self.frame_running.show_all()
            t = threading.Thread(target=self.plink.run_plink, kwargs=plinkdic)
            t.start()
            self.progressbar_plink.set_text("Searching for Runs of Homozygosity...")
            while t.isAlive():
                self.progressbar_plink.pulse()
                while gtk.events_pending():
                    gtk.main_iteration(False)
                time.sleep(0.1)
            
            # Prepare plots of the analysis results 
            t = threading.Thread(target=self.plink.mpl_plot)
            t.start()
            self.progressbar_plink.set_text("Plotting...")
            while t.isAlive():
                if self.plink.chromosome == 0:
                    self.progressbar_plink.pulse()
                else:
                    self.progressbar_plink.set_text("Plotting chromosome %i..." % self.plink.chromosome)
                    self.progressbar_plink.set_fraction(self.plink.chromosome / 26.0)
                while gtk.events_pending():
                    gtk.main_iteration(False)
                time.sleep(0.1)
            # Get a list of the images produced
            self.images = self.plink.images
            # Hide progress window
            self.frame_running.hide_all()
            # if all went well, show graphs and legend
            if len(self.images) > 0:
                self.hscale_chrom.show()
                self.hscale_chrom.set_range(0,len(self.images)-1)
                self.hscale_chrom.set_value(0)
                self.image_chrom.set_from_file(self.images[
                    int(self.hscale_chrom.get_value())])
                # if the image is taller than 900 pixels, make room for a horizontal scrollbar
                if self.image_chrom.size_request()[1] > 900:
                    self.scrolledwindow_image.set_size_request(918,900)
                self.scrolledwindow_image.show_all()
                self.frame_legend.show()
                # enable save functionality
                self.button_save.set_sensitive(True)
            else:
                # we don't want to save if there's no data
                self.button_save.set_sensitive(False)
                # likewise, hide image, slider, and legend if this is not the first run
                self.scrolledwindow_image.hide_all()
                self.hscale_chrom.hide()
                raise PlinkParamError(
                    "Something went wrong when executing plink: no images produced!")
        except PlinkParamError, err:
            self.show_error_dlg(err)
        
    def on_hscale_chrom_value_changed(self, widget):
        """Change the picture shown when the slider below it is moved."""
        self.image_chrom.set_from_file(self.images[int(widget.get_value())])
     
    def on_button_save_clicked(self, widget):
        """Save the latest results files in a tar.gz in current directory."""
        t = time.localtime()
        basename = "homruns_%d-%02d-%02d_%02d%02d%02d" % \
                (t.tm_year, t.tm_mon, t.tm_mday, t.tm_hour, t.tm_min, t.tm_sec)
        archive = tarfile.open(basename + ".tar.gz", "w:gz")
        archive.add(self.plink.workdir, basename)
        archive.close()
        self.show_msg_dlg("Results saved as %s." % (basename + ".tar.gz"))

    
class PlinkLink:
    """The class calling plink and returning its output."""
    def __init__(self):
        """Set up a temp directory for plink output"""
        self.workdir = tempfile.mkdtemp(prefix = 'homruns')
        self.images = []

    def cleanup(self):
        """Remove files created by plink and plots"""
        if os.path.exists(self.workdir):
            for f in os.listdir(self.workdir):
                os.remove(self.workdir + '/' + f)
            os.rmdir(self.workdir)
    
    def run_plink_dic(self, dic):
        return self.run_plink(**dic)

    def run_plink(self, file = None, bfile = None,
            window_kb = None, window_snp = None, window_het = None,
            window_missing = None, window_threshold = None, snp = None,
            kb = None, density = None, gap = None):
        """Run the plink genetic analysis program with the provided parameters.
        
        Requires the command plink to be in the system path."""
        # Initialize the command line. This assumes we have plink in path.
        self.plinkCommand = ['plink --noweb', '--homozyg']
        
        # Append all parameters we got to the command line.
        self.file = file
        self.bfile = bfile
        if not file is None:
            self.plinkCommand.append('--file ' + file)
        elif not bfile is None:
            self.plinkCommand.append('--bfile ' + bfile)
        self.out = self.workdir + '/plink'
        self.plinkCommand.append('--out ' + self.out)
        if not window_kb is None:
            self.plinkCommand.append('--homozyg-window-kb %.0f' % window_kb)
        if not window_snp is None:
            self.plinkCommand.append('--homozyg-window-snp %.0f' % window_snp)
        if not window_het is None:
            self.plinkCommand.append('--homozyg-window-het %.0f' % window_het)
        if not window_missing is None:
            self.plinkCommand.append('--homozyg-window-missing %.0f' % window_missing)
        if not window_threshold is None:
            self.plinkCommand.append('--homozyg-window-threshold ' + str(window_threshold))
        if not snp is None:
            self.plinkCommand.append('--homozyg-snp %.0f' % snp)
        if not kb is None:
            self.plinkCommand.append('--homozyg-kb %.0f' % kb)
        if not density is None:
            self.plinkCommand.append('--homozyg-density %.0f' % density)
        if not gap is None:
            self.plinkCommand.append('--homozyg-gap %.0f' % gap)
        # run the RoH scan
        status, output = commands.getstatusoutput(' '.join(self.plinkCommand))
        if status != 0: raise PlinkParamError('Plink had an abnormal exit code!\n\n'
                                              + "Command output: " + output)
        # and then a Mendel error scan (only makes sense if we have family data)
        self.plinkCommand = ['plink --noweb --mendel']
        self.plinkCommand.append('--out ' + self.out)
        if not self.file is None:
            self.plinkCommand.append('--file ' + self.file)
        elif not self.bfile is None:
            self.plinkCommand.append('--bfile ' + self.bfile)
        status, output = commands.getstatusoutput(' '.join(self.plinkCommand))
        if status != 0: raise PlinkParamError('Plink had an abnormal exit code!\n\n'
                                              + "Command output: " + output)
    
    def mpl_plot(self):
        """Use matplotlib to parse a successful RoH search and produce a
        visualization of the runs as a png image for each given chromosome."""
        # step 0: draw marker density histogram
        self.chromosome = 0
        markerpos = []
        for i in range(27):
            markerpos.append({})
        if not self.file is None:
            mapfile = open(self.file + '.map')
        else:
            mapfile = open(self.bfile + '.bim')
        for l in mapfile.readlines():
            fields = l.split()
            if len(fields) == 3:
                c, name, bp = fields
            elif len(fields) > 3:
                c, name, cm, bp = fields[:4]
            markerpos[int(c)][name] = int(bp)
        #
        # step 1: create a list of all the individuals with one or more
        # RoH and/or one or more Mendelian error
        self.images = []
        indlist = []
        rohinds = open(self.out + '.hom.indiv')
        for l in rohinds.readlines()[1:]:
            fields = l.split()
            if fields[3] != '0': # one or more RoH
                indlist.append(tuple(fields[:2])) # FID, IID
        rohinds.close()
        mendelinds = open(self.out + '.mendel')
        mendelerrs = []
        for i in range(27):
            mendelerrs.append([[],[]])
        for l in mendelinds.readlines()[1:]:
            fields = l.split()
            indlist.append(tuple(fields[:2]))
            mendelerrs[int(fields[2])][0].append(tuple(fields[:2])) # ID
            mendelerrs[int(fields[2])][1].append(markerpos[int(fields[2])][fields[3]]) # marker pos
        mendelinds.close()
        #
        # step 2:
        # sort the list from step 1 and assign new individual numbers starting at 1
        # to the individuals on the list using a dictionary
        indlist.sort()
        i = 1
        indmap = {}
        for ind in indlist:
            if not indmap.has_key(ind):
                indmap[ind] = i
                i += 1
        #
        # step 3: read and draw all the RoHs and Mendel errors!
        runsfile = open(self.out + '.hom')
        runs = []
        for i in range(1,27):
            runs.append([[],[],[],[]]) # (FID,IID), PHE, POS1, POS2
        for l in runsfile.readlines()[1:]:
            fields = l.split()
            chr = int(fields[3])
            runs[chr-1][0].append((fields[0],fields[1])) # (FID,IID)
            runs[chr-1][1].append(max(0,int(float(fields[2])))) # PHE, translating -9 to 0 for missing
            runs[chr-1][2].append(int(fields[6])) # POS1
            runs[chr-1][3].append(int(fields[7])) # POS2
        runsfile.close()
        plotfiles = []
        for chr, run in enumerate(runs):
            if len(run[0]) > 0 or len(mendelerrs[chr+1][0]) > 0:
                self.chromosome = chr + 1
                pyplot.clf()
                
                pyplot.xlabel('Position (bp)')
                pyplot.ylabel('Individual #')

                pyplot.xticks(size='x-small', color=grey60)
                pyplot.yticks(size='x-small', color=grey60)

                pyplot.subplots_adjust(left=0.1, right=0.9, bottom=0.05, top=0.95)
                pyplot.gcf().set_dpi(100)
                # uncomment below to make sure we have at least one pixel
                # of vertical space per individual:
                #pyplot.gcf().set_size_inches(9,max((1/0.9)*len(indmap.keys())/100, 9))
                pyplot.gcf().set_size_inches((9,9))
                
                # set up the two y axes
                rightaxis = pyplot.gca()
                leftaxis = rightaxis.figure.add_axes(rightaxis.get_position(True),
                        sharex = rightaxis, frameon = False)
                leftaxis.yaxis.tick_left()
                leftaxis.yaxis.set_label_position('left')
                rightaxis.yaxis.tick_right()
                rightaxis.yaxis.set_label_position('right')
                pyplot.setp(rightaxis.yaxis.get_ticklabels(), size = 'x-small', color = grey60)
                
                rightaxis.set_ylabel('Markers/Mb', rotation = '270')

                leftaxis.set_ylabel('Individual #')

                leftaxis.grid(color=grey80)
                
                histmax = 0
                if len(markerpos[chr+1].values()) > 0:
                    # draw a marker density histogram and get the value of the highest bar
                    histmax = max(rightaxis.hist(markerpos[chr+1].values(), 
                        bins = (chrlimits[chr][1]-chrlimits[chr][0])/1e6, 
                        range = chrlimits[chr], linewidth = 0.0, facecolor = grey80)[0])
                
              
                if len(mendelerrs[chr+1][0]) > 0:
                    leftaxis.scatter(mendelerrs[chr+1][1], [indmap[ind] for ind in mendelerrs[chr+1][0]],
                            s=10, c=mendelorange, linewidths=(0.0,)*len(mendelerrs[chr+1][0]), alpha=0.5)
                
                if len(run[0]) > 0:
                    leftaxis.hlines([indmap[ind] for ind in run[0]], # Y axis coordinate (ID)
                        run[2], run[3], # Start and end of RoH
                        colors=[(niceblue, nicegreen, nicered)[p] for p in run[1]]) # Color by phenotype


                pyplot.title('Chromosome %i' % (chr+1))
                
                rightaxis.set_xlim(chrlimits[chr])
                rightaxis.set_ylim((0, histmax * 5))
                leftaxis.set_xlim(chrlimits[chr])
                leftaxis.set_ylim((0, len(indmap.keys())))
                
                pyplot.xticks(size='x-small', color=grey60)
                pyplot.yticks(size='x-small', color=grey60)

                plotfile = self.workdir + '/mendel_homoz_chr%0i' % (chr+1)
                pyplot.savefig(plotfile)
                plotfiles.append(plotfile + '.png')
        # 
        # finally: send out a list of images (filenames)
        self.images = plotfiles
        pass
    
def look_for_file(filename, errclass):
    """Raise an errclass exception if filename is not found"""
    if not os.access(filename, os.R_OK):
        raise errclass(
        "<big><b>File missing!</b></big>\nCould not find <i>" + filename +
        "</i>, please make sure it exists")

def new_coloured_pixbuf(colour, width, height):
    """Return a (width x height) gtk.gdk.Pixbuf filled with colour.
    
    Useful when making a legend for graphs."""
    pb = gtk.gdk.Pixbuf(gtk.gdk.COLORSPACE_RGB, False, 8, width, height)
    pb.fill(colour)
    return pb

class PlinkParamError(Exception): pass

# main program starts here:
if __name__ == "__main__":
    homruns = MainWindow()
    gtk.main()
    homruns.plink.cleanup()
